import torch
import json
import os
import logging
from datetime import datetime


def setup_device():
    """
    设置计算设备
    
    Returns:
        torch.device: 计算设备
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'使用设备: {device}')
    return device


def setup_logging(log_dir, model_name):
    """
    设置日志记录
    
    Args:
        log_dir (str): 日志目录
        model_name (str): 模型名称
    
    Returns:
        str: 日志文件路径
    """
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = os.path.join(log_dir, f'{model_name}_{timestamp}.log')
    os.makedirs(log_dir, exist_ok=True)
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler()
        ]
    )
    
    return log_file


def load_label_mapping(label_file):
    """
    加载标签映射文件
    
    Args:
        label_file (str): 标签文件路径
    
    Returns:
        dict: 标签映射字典
    """
    try:
        with open(label_file, 'r', encoding='utf-8') as f:
            label_mapping = json.load(f)
        return label_mapping
    except FileNotFoundError:
        print(f"警告: 标签文件 {label_file} 未找到，将使用默认标签")
        return {}


def calculate_accuracy(outputs, targets):
    """
    计算准确率
    
    Args:
        outputs (torch.Tensor): 模型输出
        targets (torch.Tensor): 真实标签
    
    Returns:
        float: 准确率
    """
    _, predicted = torch.max(outputs.data, 1)
    total = targets.size(0)
    correct = (predicted == targets).sum().item()
    return correct / total


def save_model(model, save_path, epoch, val_acc, is_best=False):
    """
    保存模型
    
    Args:
        model (torch.nn.Module): 模型
        save_path (str): 保存路径
        epoch (int): 当前epoch
        val_acc (float): 验证准确率
        is_best (bool): 是否为最佳模型
    """
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'val_acc': val_acc,
    }
    
    torch.save(checkpoint, save_path)
    
    if is_best:
        best_path = save_path.replace('.pth', '_best.pth')
        torch.save(checkpoint, best_path)


def load_model(model, model_path, device):
    """
    加载模型
    
    Args:
        model (torch.nn.Module): 模型
        model_path (str): 模型路径
        device (torch.device): 计算设备
    
    Returns:
        torch.nn.Module: 加载后的模型
    """
    if os.path.exists(model_path):
        checkpoint = torch.load(model_path, map_location=device)
        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
            print(f"加载模型成功，epoch: {checkpoint.get('epoch', 'unknown')}, "
                  f"验证准确率: {checkpoint.get('val_acc', 'unknown'):.2f}%")
        else:
            # 兼容旧格式
            model.load_state_dict(checkpoint)
            print("加载模型成功（旧格式）")
    else:
        print(f"模型文件 {model_path} 不存在")
    
    return model


def count_parameters(model):
    """
    统计模型参数数量
    
    Args:
        model (torch.nn.Module): 模型
    
    Returns:
        int: 参数数量
    """
    return sum(p.numel() for p in model.parameters() if p.requires_grad) 